// Accelerated Dual Decomposition - An implementation of a MAP inference
// algorithm. See ai.stanford.edu/~vjojic/fastdd.pd
//
// Copyright 2011 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: nadavs@google.com (Nadav Samet)
//          noam.levy1@stanford.edu (Noam Levy)
//
// Wraps AcceleratedDD as an MRF

#ifndef _ACCELERATED_DD_MRF_H
#define _ACCELERATED_DD_MRF_H
#include "AcceleratedDD.h"
#include "mrf.h"

#include <vector>

#include <iostream>

using namespace std;

class AcceleratedDD_MRF : public MRF {
 public:
  AcceleratedDD_MRF(int width, int height, int nLabels, EnergyFunction *eng,
                    double epsilon,
                    const string &logfile = "",
                    AcceleratedDD::LOGLEVEL loglevel = AcceleratedDD::INFO);

  AcceleratedDD_MRF(int variables, int nLabels, EnergyFunction *eng,
                    double epsilon,
                    const string &logfile = "",
                    AcceleratedDD::LOGLEVEL loglevel = AcceleratedDD::INFO);
  virtual ~AcceleratedDD_MRF();

  // returns the data part of the energy
  virtual EnergyVal dataEnergy();

  // returns the smoothness part of the energy
  virtual EnergyVal smoothnessEnergy();

  virtual double lowerBound();

  // unimplemented
  virtual void setNeighbors(int pix1, int pix2, CostVal weight);

  virtual void optimizeAlg(int nIterations);

  // Returns pointer to array of size nPixels. Client may then read/write
  // solution (but not deallocate array).
  virtual Label* getAnswerPtr();
  // returns the label of the input pixel
  virtual Label getLabel(int pixel);
  // sets label of a pixel
  virtual void setLabel(int pixel,Label label);
  // sets all the labels to zero
  virtual void clearAnswer();

  // unimplemented
  virtual void setParameters(int numParam, void *param);

 protected:
  virtual void setData(DataCostFn dcost); 
  virtual void setData(CostVal* data);
  virtual void setSmoothness(CostVal* V);
  virtual void setSmoothness(SmoothCostGeneralFn cost);
  virtual void setSmoothness(int smoothExp,CostVal smoothMax, CostVal lambda);
  virtual void setCues(CostVal* hCue, CostVal* vCue); 
  virtual void initializeAlg();

 private:
  void CommonConstructor(int variables, int nLabels, const string &logfile,
                         AcceleratedDD::LOGLEVEL loglevel);

  void CheckGrid();
  void CheckNonGrid();

  MRF::EnergyVal sumFactors(const vector<AcceleratedDD::Factor*> &factors) const;

  void InnerAddGridPairwiseFactors(
      MRF::SmoothCostGeneralFn cost,
      int width,
      int x_max,
      int y_max,
      int p2_offset);
  void InnerAddEdgesFactors(MRF::SmoothCostGeneralFn cost);

  AcceleratedDD::OptimizationProblem problem_;
  double mu;
  double L;
  double epsilon_;
  AcceleratedDD::AlgorithmState *state_;
  MRF::Label *answer_;
  vector<AcceleratedDD::Factor*> dataFactors_;
  vector<AcceleratedDD::Factor*> pairwiseFactors_;
  vector<pair<AcceleratedDD::VariableId, AcceleratedDD::VariableId> > edges_;
  ostream *logstream;
};

#endif
